import torch
import os

def analyze_model_files():
    """分析模型文件大小和结构"""
    # 切换到scripts/plw目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    plw_dir = os.path.join(current_dir, 'scripts', 'plw')
    os.chdir(plw_dir)
    
    model_files = [
        'plw_model.pth',
        'plw_lstm_model.pth',
        'plw_sequence_lstm_model.pth',
        'plw_transformer_model.pth',
        'enhanced_lstm_crf_model.pth'
    ]
    
    print("模型文件分析报告")
    print("=" * 50)
    
    for model_file in model_files:
        if os.path.exists(model_file):
            # 文件大小
            size = os.path.getsize(model_file) / 1024  # KB
            print(f"{model_file}: {size:.1f}KB")
            
            # 加载模型检查点
            try:
                checkpoint = torch.load(model_file, map_location='cpu')
                print(f"  键: {list(checkpoint.keys())}")
                
                # 如果有模型状态字典，计算参数数量
                if 'model_state_dict' in checkpoint:
                    state_dict = checkpoint['model_state_dict']
                    total_params = sum(p.numel() for p in state_dict.values())
                    print(f"  参数数量: {total_params:,}")
                elif 'model' in checkpoint:
                    state_dict = checkpoint['model']
                    total_params = sum(p.numel() for p in state_dict.values())
                    print(f"  参数数量: {total_params:,}")
                elif 'lstm_model' in checkpoint:
                    state_dict = checkpoint['lstm_model']
                    total_params = sum(p.numel() for p in state_dict.values())
                    print(f"  参数数量: {total_params:,}")
                elif 'ensemble_config' in checkpoint:
                    print(f"  集成模型配置: {checkpoint['ensemble_config']}")
                    # 集成模型不直接计算参数数量
                    total_params = 0
                    if 'lstm_model' in checkpoint:
                        total_params += sum(p.numel() for p in checkpoint['lstm_model'].values())
                    if 'transformer_model' in checkpoint:
                        total_params += sum(p.numel() for p in checkpoint['transformer_model'].values())
                    print(f"  参数数量(估计): {total_params:,}")
                else:
                    print(f"  无法确定参数数量")
                    
            except Exception as e:
                print(f"  加载失败: {e}")
        else:
            print(f"{model_file}: 文件不存在")
        print()

if __name__ == "__main__":
    analyze_model_files()